class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def isBalanced(self, root):
        """
        :type root: TreeNode
        :rtype: bool
        """

        def count(node):
            if not node:
                return True, 0
            else:
                left_flag, left_count = count(node.left)
                right_flag, right_count = count(node.right)
                if left_flag and right_flag:
                    if -1 <= right_count - left_count <= 1:
                        return True, max(left_count, right_count) + 1
                    else:
                        return False, 0
                else:
                    return False, 0

        flag, _ = count(root)

        return flag
